import asyncio
import os
from pyexpat.errors import messages
from typing import Optional

from mcp import ClientSession
from openai import OpenAI
from dotenv import load_dotenv
from contextlib import  AsyncExitStack

load_dotenv()

class MCPClient:
    def __init__(self):
        self.exit_stack=AsyncExitStack()
        self.openai_api_key = os.getenv("OPENAI_API_KEY")
        self.base_url = os.getenv("BASE_URL")
        self.mode = os.getenv("MODEL")
        if not self.openai_api_key:
            raise ValueError("未找到openai api key")
        print("---",self.mode)
        self.client=OpenAI(api_key=self.openai_api_key,base_url=self.base_url)
        self.session:Optional[ClientSession]=None
    async def connect_to_server(self,server_script_path):
        is_python=server_script_path.endswith(".py")
        is_js=server_script_path.endswith(".js")
        if not is_python and not is_js:
            raise ValueError("只支持python和js脚本")
        command="python" if is_python else "node"
        server_params=["--no-warnings"] if is_js else []
    async def process_query(self,query:str)->str:
        messages=[{"role":"user","content":query},{
            "role":"system",
            "content":"你是一个知识库助手，回答问题时要尽量引用知识库中的内容"
        }]
        try:
            response=await asyncio.get_event_loop().run_in_executor(
                None,
                lambda :self.client.chat.completions.create(
                    model=self.mode,
                    messages=messages
                )
            )
            return response.choices[0].message.content
        except Exception as e:
            return f"调用openai api失败: {str(e)}"
    async def chat_loop(self):
        while True:
            try:
                query = input("请输入问题：").strip()
                if query.lower() == "exit":
                    break
                response = await self.process_query(query)
                print(f"回答：{response}")
            except Exception as e:
                print(f"发生错误，请重试:{str(e)}")

    async def clean_up(self):
        await self.exit_stack.aclose()

async def main():
    client=MCPClient()
    try:
        await client.chat_loop()
    finally:
        await client.clean_up()
if __name__=="__main__":
    asyncio.run(main())

